# -*- coding: utf-8 -*-

import pydicom
import os
import numpy as np
import SimpleITK as sitk


def get_tiff_images(dcm_folder, dcm_number):
    global dcm_dict
    onePatien = []
    zLocation = []
    for i in range(dcm_number):
        if i in range(0, 10):
            file_name = '00000' + str(i) + '.dcm'
        if i in range(10, 100):
            file_name = '0000' + str(i) + '.dcm'
        if i in range(100, dcm_number):
            file_name = '000' + str(i) + '.dcm'
        ds = pydicom.dcmread(os.path.join(dcm_folder, file_name))
        window_width = int(ds.WindowWidth)
        window_center = int(ds.WindowCenter)
        window_min = window_center - 0.5 * window_width
        window_max = window_center + 0.5 * window_width
        UID = ds.SOPInstanceUID
        Z_pos = ds.ImagePositionPatient[2]
        dcm_dict[UID] = os.path.join(dcm_folder, file_name)
        image = sitk.ReadImage(os.path.join(dcm_folder, file_name))
        image_array = sitk.GetArrayFromImage(image)
        ct_data = image_array.reshape(512, 512)

        ct_data[ct_data < window_min] = window_min
        ct_data[ct_data > window_max] = window_max

        ct_data = np.uint8((ct_data - window_min) / (window_max - window_min) * 255)
        onePatien.append(ct_data)
        zLocation.append(Z_pos)
    index = list(range(len(zLocation)))
    index.sort(key=lambda s: -zLocation[s])
    onePatien = np.array(onePatien)[index]
    print(len(dcm_dict))
    return dcm_dict, onePatien



def returnPatients(root_path):
    global dcm_dict
    dcm_dict = {}
    patients = []
    for root, dirs, files in os.walk(root_path):
        if len(dirs) == 0:

            f_root = os.path.dirname(root).split('/')[-1]
            print(f_root)

            if len(files) > 1:

                if '.DS_Store' in files:
                    dcm_dict, onePatient = get_tiff_images(root, len(files) - 1)
                else:
                    dcm_dict, onePatient = get_tiff_images(root, len(files))
                patients.append(onePatient)
    return patients


def preprocess(path):

    rawData = returnPatients(path)
    numOfPatiens = len(rawData)
    for i in range(numOfPatiens):
        train_x = rawData[i].copy()
        location = np.zeros([train_x.shape[0], train_x.shape[1], train_x.shape[2]] + [3], dtype='float32')
        for m in range(train_x.shape[0]):
            for n in range(train_x.shape[1]):
                 for k in range(train_x.shape[2]):
                    location[m, n, k][0] = m * (150/(train_x.shape[1]))
                    location[m, n, k][1] = n
                    location[m, n, k][2] = k
                    # print(location[m,j,k])
        train_x = train_x[:, np.newaxis, :, :]
        location = location.transpose(0, 3, 1, 2)
        new_train_x = np.concatenate((train_x, location), 1)
        return new_train_x
